package com.briup.mock.init;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.briup.mock.properties.ZookeeperProperties;
import com.briup.mock.result.MockResult;
import com.briup.mock.result.ZookeeperEvaluateResult;
import com.briup.mock.result.ZookeeperInitResult;
import com.briup.mock.result.ZookeeperOptimizeResult;
import com.briup.mock.vo.Evaluate;
import com.briup.mock.vo.EvaluateLabel;
import com.briup.mock.vo.Training;
import com.briup.mock.vo.TrainingLabel;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.recipes.cache.TreeCache;
import org.apache.curator.framework.recipes.cache.TreeCacheEvent;
import org.apache.zookeeper.data.Stat;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.CommandLineRunner;
import org.springframework.boot.env.OriginTrackedMapPropertySource;
import org.springframework.core.env.Environment;
import org.springframework.stereotype.Component;
import org.springframework.web.context.support.StandardServletEnvironment;

import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 *
 * springboot项目启动后，进行自定义的初始化操作
 * @author wood
 */
@Component
public class ApplicationInit implements CommandLineRunner {

    @Value("${server.port}")
    private int port;
    @Autowired
    private Environment env;
    //正则匹配规则
    private static final String REG;
    //正则表达式对象
    private static final Pattern PATTERN;

    static {
        REG = "\\[(.*)\\]";
        PATTERN = Pattern.compile(REG);
    }

    @Autowired
    private CuratorFramework client;

    @Autowired
    private ZookeeperProperties zookeeperProperties;

    @Override
    public void run(String... args) throws Exception {
        // 启动后显示基本信息
        showInfo();

        // 启动后对zookeeper的操作
        zookeeperOps();
    }

    // =========================面是zookeeperOps相关方法===========================
    private void zookeeperOps() throws Exception{
        // 检查节点是否存在，不存在则创建
        checkNode(zookeeperProperties.getTaskPath());
        checkNode(zookeeperProperties.getResultPath());

        // 读取节点上的数据，并处理（启动后的第一次处理，防止对之前数据的遗漏）
        byte[] bytes = client.getData().forPath(zookeeperProperties.getTaskPath());
        // 如果数据是null或者不是{开头，那么就退出
        this.processZKData(bytes);

        // 获取TreeCache对象，准备给节点添加监听器
        TreeCache treeCache = new TreeCache(client, zookeeperProperties.getTaskPath());

        // 添加监听器
        treeCache.getListenable().addListener((cli, event) -> {
            // 如果指定节点有更新，就进行处理
            if (event.getType() == TreeCacheEvent.Type.NODE_UPDATED) {

                // 如果触发事件的节点不是指的节点，那么就方法退出
                if (!event.getData().getPath().equals(zookeeperProperties.getTaskPath())) {
                    return;
                }
                // 获取指定节点上的数据
                byte[] data = event.getData().getData();
                // 处理数据
                this.processZKData(data);
            }

        });

        treeCache.start();
    }

    /**
     * 处理zk事件获取的节点上的数据
     * @param data  节点数据
     * @throws Exception 异常
     */
    private void processZKData(byte[] data)throws Exception{
        // 如果数据是null或者不是{开头，那么就退出
        if(data==null || !new String(data).startsWith("{")){
            return;
        }

        String jsonStr = new String(data);
        JSONObject jsonObject = JSON.parseObject(jsonStr);
        // 判断是哪种任务
        Integer taskType = (Integer) jsonObject.get("taskType");

        ObjectMapper objectMapper = new ObjectMapper();
        if (taskType != null) {
            switch (taskType) {
                case 0:
                    // 初始化训练
                    ZookeeperInitResult zookeeperResult = objectMapper.readValue(jsonStr, ZookeeperInitResult.class);
                    Thread.sleep(3000);
                    // 返回处理结果
                    MockResult mockResult = mockInitResult(zookeeperResult);
                    String mockResultJson = objectMapper.writeValueAsString(mockResult);
                    client.setData().forPath(zookeeperProperties.getResultPath(), mockResultJson.getBytes());
                    break;
                case 1:
                    // 初优化训练
                    ZookeeperOptimizeResult zookeeperOptimizeResult = objectMapper.readValue(jsonStr, ZookeeperOptimizeResult.class);
                    Thread.sleep(3000);
                    // 返回处理结果
                    MockResult optResult = mocOptimizeResult(zookeeperOptimizeResult);
                    String optResultJson = objectMapper.writeValueAsString(optResult);
                    client.setData().forPath(zookeeperProperties.getResultPath(), optResultJson.getBytes());
                    break;
                case 2:
                    // 评估训练
                    ZookeeperEvaluateResult zookeeperEvaluateResult = objectMapper.readValue(jsonStr, ZookeeperEvaluateResult.class);
                    Thread.sleep(3000);
                    // 返回处理结果
                    MockResult evaluateResult = mockEvaluateResult(zookeeperEvaluateResult);
                    String evluateResultJson = objectMapper.writeValueAsString(evaluateResult);
                    client.setData().forPath(zookeeperProperties.getResultPath(), evluateResultJson.getBytes());
                    break;
                default:

            }
        }

        // 数据处理完成后，把Task节点的值设置为null
        client.setData().forPath(zookeeperProperties.getTaskPath(),null);

        Thread.sleep(3000);
    }

    // 模拟评估的结果
    private MockResult mockEvaluateResult(ZookeeperEvaluateResult zookeeperEvaluateResult) {
        Integer trainingId = zookeeperEvaluateResult.getTrainingId();
        Integer datasestId = zookeeperEvaluateResult.getDatasestId();
        List<String> classNameList = zookeeperEvaluateResult.getClassNameList();
        // 封装评估信息
        Evaluate evaluate = new Evaluate();
        evaluate.setAccuracyRate(0.7F);
        evaluate.setEvaluateTime(new Date());
        evaluate.setTrainingId(trainingId);
        evaluate.setDatasetId(datasestId);
        // 封装评估分类信息
        List<EvaluateLabel> evaluateLabelList = new ArrayList<>();
        classNameList.forEach(className->{
            EvaluateLabel evaluateLabel = new EvaluateLabel();
            evaluateLabel.setLabelTime(new Date());
            evaluateLabel.setLabelName(className);
            evaluateLabel.setLabelDescribe(className + "测试描述");
            evaluateLabel.setF1score(0.6F);
            evaluateLabel.setGscore(0.6F);
            evaluateLabel.setPrecisionRate(0.6F);
            evaluateLabel.setRecallRate(0.6F);
            evaluateLabelList.add(evaluateLabel);
        });
        // 封装评估错误信息
        MockResult mockResult = new MockResult();
        mockResult.setEvaluate(evaluate);
        mockResult.setEvaluateLabelList(evaluateLabelList);
        //mockResult.setEvaluateErrorList();

        mockResult.setResultType(2);
        return mockResult;
    }

    // 模拟优化训练结果数据
    private MockResult mocOptimizeResult(ZookeeperOptimizeResult zookeeperOptimizeResult) {
        Training training = new Training();
        Set<String> labelNameSet = zookeeperOptimizeResult.getLabelNameSet();
        List<Integer> datasetIdList = zookeeperOptimizeResult.getDatasetIdList();
        String lastVersion = zookeeperOptimizeResult.getLastVersion();
        MockResult mockResult = new MockResult();
        // 封装训练信息
        training.setTrainingTime(new Date());
        training.setAccuracyRate(0.5F);
        // 更换版本
        training.setModelVersion(lastVersion.substring(0, 1) + (Integer.parseInt(lastVersion.substring(1)) + 1));
        training.setModelFileaddr(zookeeperOptimizeResult.getNewModelPath());
        training.setModelId(zookeeperOptimizeResult.getModelId());
        mockResult.setTraining(training);
        // 封装训练分类信息
        List<TrainingLabel> trainingLabelList = new ArrayList<>();
        labelNameSet.forEach(labelName -> {
            TrainingLabel trainingLabel = new TrainingLabel();
            trainingLabel.setLabelTime(new Date());
            trainingLabel.setLabelName(labelName);
            trainingLabel.setLabelDescribe(labelName + "测试描述");
            trainingLabel.setF1score(0.6F);
            trainingLabel.setGscore(0.6F);
            trainingLabel.setPrecisionRate(0.6F);
            trainingLabel.setRecallRate(0.6F);
            trainingLabelList.add(trainingLabel);
        });
        mockResult.setTrainingLabelList(trainingLabelList);
        // 封装训练需要用到数据集
        mockResult.setDatasetIdList(datasetIdList);
        // 结果类型 1 优化训练
        mockResult.setResultType(1);
        return mockResult;

    }

    // 模拟初始化结果数据
    private MockResult mockInitResult(ZookeeperInitResult zookeeperResult) {
//        Training training = zookeeperResult.getTraining();
        Training training = new Training();
        Set<String> labelNameSet = zookeeperResult.getLabelNameSet();
        List<Integer> datasetIdList = zookeeperResult.getDatasetIdList();

        MockResult mockResult = new MockResult();
        // 封装训练信息
        training.setTrainingTime(new Date());
        training.setAccuracyRate(0.5F);
        training.setModelVersion("V0");
        // 第一此训练会传默认的模型地址
        training.setModelFileaddr(zookeeperResult.getNewModelPath());
        training.setModelId(zookeeperResult.getModelId());
        mockResult.setTraining(training);
        // 封装训练分类信息
        List<TrainingLabel> trainingLabelList = new ArrayList<>();
        labelNameSet.forEach(labelName -> {
            TrainingLabel trainingLabel = new TrainingLabel();
            trainingLabel.setLabelTime(new Date());
            trainingLabel.setLabelName(labelName);
            trainingLabel.setLabelDescribe(labelName + "测试描述");
            trainingLabel.setF1score(0.5F);
            trainingLabel.setGscore(0.5F);
            trainingLabel.setPrecisionRate(0.5F);
            trainingLabel.setRecallRate(0.5F);
            trainingLabelList.add(trainingLabel);
        });
        mockResult.setTrainingLabelList(trainingLabelList);
        // 封装训练需要用到数据集
        mockResult.setDatasetIdList(datasetIdList);
        // 结果类型 0 初始化训练
        mockResult.setResultType(0);
        return mockResult;
    }

    /**
     * 检查节点是否存在，不存在则创建
     * @param nodePath 节点路径
     * @throws Exception 异常
     */
    private void checkNode(String nodePath) throws Exception {
        Stat stat = null;
        stat = client.checkExists().forPath(nodePath);
        if (stat == null) {
            client.create().creatingParentsIfNeeded().forPath(nodePath);
        }
        stat = client.checkExists().forPath(nodePath);
        if (stat == null) {
            client.create().creatingParentsIfNeeded().forPath(nodePath);
        }
    }





    // =========================下面是showInfo相关方法===========================
    private void showInfo(){

        System.out.println("*****************本次启动激活的profile*****************");
        System.out.println("当前激活的profile: "+ Arrays.toString(env.getActiveProfiles()));
        System.out.println("*****************本次启动激活的profile*****************");

        System.out.println();

        System.out.println("*****************本次启动读取的配置文件*****************");
        StandardServletEnvironment sse = (StandardServletEnvironment) env;
        sse.getPropertySources().forEach(ps->{
            if(ps instanceof OriginTrackedMapPropertySource){
                OriginTrackedMapPropertySource ot = (OriginTrackedMapPropertySource) ps;
                System.out.println(findConfigFile(ot.getName()));
            }
        });
        System.out.println("*****************本次启动读取的配置文件*****************");
    }

    /**
     * 使用正则表达式，查找配置文件，例如 [application.yml]
     * @param line
     * @return
     */
    private static String findConfigFile(String line) {

        Matcher m = PATTERN.matcher(line);

        if(m.find()){
            return m.group();
        }

        return null;

    }

}
